#include "FastLED.h" 
#include <Arduino.h>

#define BAUDRATE 115200l
#define COMM_HEAD '#'
#define COMM_TAIL '!'
#define COMM_FILTER ' '
#define COMM_STATE_WAITFORHEAD 0
#define COMM_STATE_WAITFORTAIL 1
#define COMM_MAX_BUFFER_SIZE 256

#define MOTOR_PID_KP 0.1f
#define MOTOR_PID_KI 0.02f
#define MOTOR_MAX_PWM 255
#define MOTOR_LEFT 0
#define MOTOR_RIGHT 1

#define DATA_PIN 8       //设置LED引脚
#define LED_TYPE WS2812  //灯带类型
#define NUM_LEDS 50      //灯珠数量
#define COLOR_ORDER GRB  //灯珠顺序
#define BRIGHTNESS 128
uint8_t max_brightness = 128;  //最大亮度

CRGB leds[NUM_LEDS];  // 给灯带命名

#define DIN_PIN 5       //测量电压
#define BUTTON_PIN 9    //按钮输入

// motor pins
// A - left B - right
const int EN = 12;
const int AIN1 = 15;
const int AIN2 = 14;
const int PWMA = 13;
const int BIN1 = 16;
const int BIN2 = 17;
const int PWMB = 18;
const int ENCAA = 1;
const int ENCAB = 2;
const int ENCBA = 3;
const int ENCBB = 4;


// motor encoders data
long encoderLeftPeriodCount = 0;
long encoderRightPeriodCount = 0;
long encoderLeftCount = 0;
long encoderRightCount = 0;

// speed control
int motorLeftTargetSpeedPPS = 0;
int motorRightTargetSpeedPPS = 0;
int motorLeftCurrentSpeedPPS = 0;
int motorRightCurrentSpeedPPS = 0;

// control freq
long currentMillis = 0;
long previousMillis = 0;

// pid params
int lastErrorLeft = 0;
int lastErrorRight = 0;
float PWMOutputLeft = 0;
float PWMOutputRight = 0;

// comm
int comm_state = COMM_STATE_WAITFORHEAD;
String comm_packet_buffer;

// light
int currentLightColor = -1;

// battery
int batteryCount = 0;
double batterySum = 0;
double batteryMeasure = 0;
double batteryValue = 11.6;

// button
int buttonState = 0;
int buttonLastMillis = 0;

// motor A encoder ISP
void ISR_EncoderAA() {
  if (digitalRead(ENCAA) == HIGH) {
    if (digitalRead(ENCAB) == LOW) {
      encoderLeftPeriodCount++;
    } else {
      encoderLeftPeriodCount--;
    }
  } else {
    if (digitalRead(ENCAB) == HIGH) {
      encoderLeftPeriodCount++;
    } else {
      encoderLeftPeriodCount--;
    }
  }
}

void ISR_EncoderAB() {
  if (digitalRead(ENCAB) == HIGH) {
    if (digitalRead(ENCAA) == HIGH) {
      encoderLeftPeriodCount++;
    } else {
      encoderLeftPeriodCount--;
    }
  } else {
    if (digitalRead(ENCAA) == LOW) {
      encoderLeftPeriodCount++;
    } else {
      encoderLeftPeriodCount--;
    }
  }
}

// motor B encoder ISP
void ISR_EncoderBA() {
  if (digitalRead(ENCBA) == HIGH) {
    if (digitalRead(ENCBB) == LOW) {
      encoderRightPeriodCount++;
    } else {
      encoderRightPeriodCount--;
    }
  } else {
    if (digitalRead(ENCBB) == HIGH) {
      encoderRightPeriodCount++;
    } else {
      encoderRightPeriodCount--;
    }
  }
}

void ISR_EncoderBB() {
  if (digitalRead(ENCBB) == HIGH) {
    if (digitalRead(ENCBA) == HIGH) {
      encoderRightPeriodCount++;
    } else {
      encoderRightPeriodCount--;
    }
  } else {
    if (digitalRead(ENCBA) == LOW) {
      encoderRightPeriodCount++;
    } else {
      encoderRightPeriodCount--;
    }
  }
}

int GetEncoderData(int motorID) {
  int retValue = 0;

  if (motorID == MOTOR_LEFT) {
    retValue = encoderLeftPeriodCount;
    encoderLeftPeriodCount = 0;
  } else if (motorID == MOTOR_RIGHT) {
    retValue = encoderRightPeriodCount;
    encoderRightPeriodCount = 0;
  }

  return retValue;
}

void MotorRightSetPWM(int pwmValue) {
  // direction
  if (pwmValue < 0) {
    digitalWrite(AIN1, LOW);
    digitalWrite(AIN2, HIGH);
    pwmValue = pwmValue < -MOTOR_MAX_PWM ? -MOTOR_MAX_PWM : pwmValue;
  } else if (pwmValue > 0) {
    digitalWrite(AIN1, HIGH);
    digitalWrite(AIN2, LOW);
    pwmValue = pwmValue > MOTOR_MAX_PWM ? MOTOR_MAX_PWM : pwmValue;
  } else {
    digitalWrite(AIN1, LOW);
    digitalWrite(AIN2, LOW);
  }

  // pwm
  analogWrite(PWMA, abs(pwmValue));
}

void MotorLeftSetPWM(int pwmValue) {
  // direction
  if (pwmValue < 0) {
    digitalWrite(BIN1, LOW);
    digitalWrite(BIN2, HIGH);
    pwmValue = pwmValue < -MOTOR_MAX_PWM ? -MOTOR_MAX_PWM : pwmValue;
  } else if (pwmValue > 0) {
    digitalWrite(BIN1, HIGH);
    digitalWrite(BIN2, LOW);
    pwmValue = pwmValue > MOTOR_MAX_PWM ? MOTOR_MAX_PWM : pwmValue;
  } else {
    digitalWrite(BIN1, LOW);
    digitalWrite(BIN2, LOW);
  }

  // pwm
  analogWrite(PWMB, abs(pwmValue));
}

int MotorLeftPIDCalc(int currentSpeedPPS, int targetSpeedPPS) {
  int currentError = targetSpeedPPS - currentSpeedPPS;
  PWMOutputLeft += MOTOR_PID_KP * (currentError - lastErrorLeft) + MOTOR_PID_KI * currentError;
  lastErrorLeft = currentError;
  PWMOutputLeft = PWMOutputLeft > MOTOR_MAX_PWM ? MOTOR_MAX_PWM : (PWMOutputLeft < -1 * MOTOR_MAX_PWM ? -1 * MOTOR_MAX_PWM : PWMOutputLeft);
  return (int)PWMOutputLeft;
}

int MotorRightPIDCalc(int currentSpeedPPS, int targetSpeedPPS) {
  int currentError = targetSpeedPPS - currentSpeedPPS;
  PWMOutputRight += MOTOR_PID_KP * (currentError - lastErrorRight) + MOTOR_PID_KI * currentError;
  lastErrorRight = currentError;
  PWMOutputRight = PWMOutputRight > MOTOR_MAX_PWM ? MOTOR_MAX_PWM : (PWMOutputRight < -1 * MOTOR_MAX_PWM ? -1 * MOTOR_MAX_PWM : PWMOutputRight);
  return (int)PWMOutputRight;
}

void LightSetColor(int color) {
  if (currentLightColor != color) {
    for (int i = 0; i < NUM_LEDS; i++) {
      switch (color) {
        // white
        case 0:
          leds[i] = CRGB(255, 255, 255);
          break;
          // red
        case 1:
          leds[i] = CRGB(255, 0, 0);
          break;
        // green
        case 2:
          leds[i] = CRGB(0, 255, 0);
          break;
        // blue
        case 3:
          leds[i] = CRGB(0, 0, 255);
          break;
      }
    }

    FastLED.show();

    currentLightColor = color;
  }
}

void setup() {
  // serial begin
  Serial.begin(BAUDRATE);

  //电压检测
  analogReadResolution(12);
  // adc.attach(DIN_PIN);
  pinMode(DIN_PIN, INPUT);

  // motor pins
  pinMode(AIN1, OUTPUT);
  pinMode(AIN2, OUTPUT);
  pinMode(BIN1, OUTPUT);
  pinMode(BIN2, OUTPUT);
  pinMode(PWMA, OUTPUT);
  pinMode(PWMB, OUTPUT);
  pinMode(EN, OUTPUT);
  pinMode(ENCAA, INPUT);
  pinMode(ENCAB, INPUT);
  pinMode(ENCBA, INPUT);
  pinMode(ENCBB, INPUT);
  pinMode(BUTTON_PIN, INPUT);

  // motor init
  digitalWrite(EN, HIGH);

  // interrupts
  attachInterrupt(digitalPinToInterrupt(ENCAA), ISR_EncoderAA, CHANGE);
  attachInterrupt(digitalPinToInterrupt(ENCAB), ISR_EncoderAB, CHANGE);
  attachInterrupt(digitalPinToInterrupt(ENCBA), ISR_EncoderBA, CHANGE);
  attachInterrupt(digitalPinToInterrupt(ENCBB), ISR_EncoderBB, CHANGE);

  // light
  FastLED.addLeds<LED_TYPE, DATA_PIN, COLOR_ORDER>(leds, NUM_LEDS);

  LightSetColor(0);

  // millis
  previousMillis = millis();
}

void loop() {
  while (Serial.available()) {
    char current_char = Serial.read();

    if (current_char == COMM_HEAD) {
      comm_packet_buffer = "";
      comm_state = COMM_STATE_WAITFORTAIL;
    } else if (comm_state == COMM_STATE_WAITFORTAIL && current_char == COMM_TAIL) {
      int dataIndex = 0;
      int lastFilterPos = 0;
      int bufferLength = comm_packet_buffer.length();

      for (int i = 0; i < bufferLength; i++) {
        if (comm_packet_buffer.charAt(i) == COMM_FILTER or i == bufferLength - 1) {
          if (dataIndex == 0) {
            motorRightTargetSpeedPPS = comm_packet_buffer.substring(lastFilterPos, i).toInt();
          } else if (dataIndex == 1) {
            motorLeftTargetSpeedPPS = comm_packet_buffer.substring(lastFilterPos + 1, i).toInt();
          } else if (dataIndex == 2) {
            LightSetColor(comm_packet_buffer.substring(lastFilterPos + 1, bufferLength).toInt());
          }

          lastFilterPos = i;
          dataIndex++;
        }
      }

      // int filterPos = comm_packet_buffer.indexOf(COMM_FILTER);

      // if (filterPos != -1) {
      //   motorLeftTargetSpeedPPS = comm_packet_buffer.substring(0, filterPos).toInt();
      //   motorRightTargetSpeedPPS = comm_packet_buffer.substring(filterPos + 1, comm_packet_buffer.length()).toInt();
      // }

      comm_packet_buffer = "";
      comm_state = COMM_STATE_WAITFORHEAD;
    } else if (comm_state == COMM_STATE_WAITFORTAIL) {
      comm_packet_buffer += current_char;

      if (comm_packet_buffer.length() >= COMM_MAX_BUFFER_SIZE) {
        comm_packet_buffer = "";
        comm_state = COMM_STATE_WAITFORHEAD;
      }
    }
  }

  currentMillis = millis();

  // 50hz
  if (currentMillis - previousMillis >= 20) {

    // button.tick();
    // get pulses in this cycle
    int encoderLeftCountThisCycle = GetEncoderData(MOTOR_LEFT);
    int encoderRightCountThisCycle = GetEncoderData(MOTOR_RIGHT);

    // calc pulse count
    encoderLeftCount += encoderLeftCountThisCycle;
    encoderRightCount += encoderRightCountThisCycle;

    // calc motor speed
    int timeDiff = currentMillis - previousMillis;

    motorLeftCurrentSpeedPPS = (int)(encoderLeftCountThisCycle * 1000.0 / timeDiff);
    motorRightCurrentSpeedPPS = (int)(encoderRightCountThisCycle * 1000.0 / timeDiff);

    // PID control
    MotorLeftSetPWM(MotorLeftPIDCalc(motorLeftCurrentSpeedPPS, motorLeftTargetSpeedPPS));
    MotorRightSetPWM(MotorRightPIDCalc(motorRightCurrentSpeedPPS, motorRightTargetSpeedPPS));

    previousMillis = currentMillis;

    if (digitalRead(BUTTON_PIN) == 1) {
      if (buttonState == 0) {
        buttonState = 1;
        buttonLastMillis = millis();
      } else if (buttonState == 1 && millis() - buttonLastMillis >= 50) {
        buttonState = 2;
      }
    } else {
      buttonState = 0;
    }

    batteryMeasure = (((analogReadMilliVolts(DIN_PIN) - 240) / 980.0) * 11);

    if ((batteryMeasure >= 10.0) && (batteryMeasure <= 13.0)) {
      batteryCount++;
      batterySum += batteryMeasure;
    }

    if (batteryCount == 20) {
      batteryValue = (batterySum / 20.0);
      batteryCount = 0;
      batterySum = 0;
    }

    // send data
    Serial.print(COMM_HEAD);
    Serial.print(motorRightCurrentSpeedPPS);
    Serial.print(COMM_FILTER);
    Serial.print(motorLeftCurrentSpeedPPS);
    Serial.print(COMM_FILTER);
    Serial.print(encoderLeftCount);
    Serial.print(COMM_FILTER);
    Serial.print(encoderRightCount);
    Serial.print(COMM_FILTER);
    Serial.print((int)((batteryValue)*10));
    Serial.print(COMM_FILTER);
    Serial.print(buttonState == 2 ? 1 : 0);
    Serial.println(COMM_TAIL);
  } else {
  }
}
